import cv2
import numpy as np
import os
from glob import glob
from tqdm import tqdm

def process_image(img_path, output_dir, prefix):
    """处理单张图片的核心函数（针对科研拼图优化）"""
    try:
        # 使用更可靠的路径处理方式
        img_path = os.path.normpath(img_path)
        
        # 尝试以不同编码方式打开文件
        img = None
        try:
            img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        except Exception as e:
            print(f"Failed to read {img_path} with first method: {str(e)}")
            try:
                img = cv2.imread(img_path)
            except Exception as e:
                print(f"Failed to read {img_path} with second method: {str(e)}")
                return 0
        
        if img is None:
            print(f"Warning: Failed to load image {img_path}")
            return 0
        
        # 预处理（增强边界检测）
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blur = cv2.medianBlur(gray, 5)
        thresh = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
                                     cv2.THRESH_BINARY_INV, 11, 7)
        # 形态学处理（连接断裂边界）
        kernel = np.ones((9,9), np.uint8)
        dilated = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
        
        # 查找轮廓
        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 筛选有效子图轮廓（针对科研拼图特征）
        valid_contours = []
        for cnt in contours:
            x,y,w,h = cv2.boundingRect(cnt)
            area = w*h
            aspect = w/float(h)
            
            # 筛选条件（根据科研图表特征调整）
            if (area > img.shape[0]*img.shape[1]*0.04 and  # 最小面积4%
                0.3 < aspect < 3.0 and                     # 更宽松的宽高比
                max(w,h) > img.shape[0]*0.2):              # 最小边长限制
                valid_contours.append((x,y,w,h))
        
        # 按科研论文常见的阅读顺序排序（从左到右，从上到下）
        valid_contours.sort(key=lambda b: (b[1]//(img.shape[0]//10), b[0]))  # 动态Y轴分组
        
        # 保存子图（使用连续字母命名）
        saved_count = 0
        for i, (x,y,w,h) in enumerate(valid_contours[:26]):  # 最多26个子图
            sub_img = img[y:y+h, x:x+w]
            letter = chr(65 + i)  # A-Z
            output_name = f"{prefix}_{letter}.png"
            output_path = os.path.join(output_dir, output_name)
            
            # 确保输出目录存在
            os.makedirs(output_dir, exist_ok=True)
            
            # 使用更可靠的保存方式
            try:
                cv2.imwrite(output_path, sub_img)
                saved_count += 1
            except Exception as e:
                print(f"Failed to save {output_path}: {str(e)}")
        
        return saved_count
    except Exception as e:
        print(f"Error processing {img_path}: {str(e)}")
        return 0


def batch_process(input_root="images", output_dir="all_subfigures"):
    """批量处理所有图片到同一输出目录"""
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有图片路径（使用更可靠的方式）
    img_paths = []
    for root, _, files in os.walk(input_root):
        for f in files:
            if f.lower().endswith(('.png','.jpg','.jpeg','.bmp','.tiff')):
                full_path = os.path.join(root, f)
                if os.path.exists(full_path):
                    img_paths.append((root, f))
                else:
                    print(f"Warning: File not found {full_path}")
    
    print(f"Found {len(img_paths)} valid images to process")
    
    total_saved = 0
    for root, fname in tqdm(img_paths, desc="Processing"):
        # 生成前缀：父文件夹名前8字符+文件名
        folder_part = os.path.basename(root)[:8]
        file_part = os.path.splitext(fname)[0]
        clean_prefix = "".join(c for c in f"{folder_part}_{file_part}" if c.isalnum())
        
        img_path = os.path.join(root, fname)
        saved = process_image(img_path, output_dir, clean_prefix)
        total_saved += saved
    
    print(f"Completed! Saved {total_saved} subfigures to {output_dir}")

if __name__ == "__main__":
    # 使用示例
    batch_process(input_root="output_images", output_dir="all_subfigures")